#!/home/zhonc0a/Enthought/Canopy_64bit/User/bin/python
import shutil
import os
import argparse
import sys
import re
import subprocess
class converge():
    def __init__(self):
        self.encut=''
        self.prec=''
        self.nomega=''
        self.nwlist=''
        self.nbands=''
        self.nblist=''
        self.eclist=''
        self.kpoint=''
        self.kplist=''
        self.nbandsgw=''
        self.nglist=''
        self.encutgw=''
        self.eglist=''
        self.baseDIR=''
        self.functional='HSE03'
        self.parameter=[]
        self.nproc='20'
        self.dirBaseName=[]
        self.currentDIR=os.getcwd()
        self.steps=''
                
#read the parmeter from the incar and kpoints file 
    def readparam(self):
        os.chdir(self.baseDIR)
        incar=open('INCAR')
        for line in incar:
            if 'ENCUTGW' in line:
                self.encutgw=line.split('=')[1].rstrip('\n')
            if 'NOMEGA' in line:
                self.nomega=line.split('=')[1].rstrip('\n')
            #if 'NBANDSGW' in line:
            #    self.nbandsgw=line.split('=')[1].rstrip('\n')
            if 'NBANDS' in line:
                self.nbands=line.split('=')[1].rstrip('\n')
            if 'PREC=' in line:
                self.prec=line.split('=')[1].rstrip('\n')
            if 'ENCUT=' in line:
                self.encut=line.split('=')[1].rstrip('\n')
        kpoint=open('KPOINTS')
        for line in kpoint:
            if re.search('^[0-9]+\s+[0-9]+\s+[0-9]+\s*$',line):
                self.kpoint=','.join(re.split('\s+',line))
        os.chdir(self.currentDIR)
    
        
        
#construct directory and input file      
    def setformation(self):
        #construct parameters for inpgen_vasp.py and dirname
        if self.kpoint:
            self.parameter=self.parameter+[' -k '+self.kpoint]
        else:
            self.parameter=self.parameter+[' ']
        self.dirBaseName=self.dirBaseName+['K'+''.join(self.kpoint.split(','))]
        if self.nbands:
            self.parameter=self.parameter+[' -b '+self.nbands]
        else:
            self.parameter=self.parameter+[' ']
        self.dirBaseName=self.dirBaseName+['B'+self.nbands]
        if self.nomega:
            self.parameter=self.parameter+[' -w '+self.nomega]
        else:
            self.parameter=self.parameter+[' ']
        self.dirBaseName=self.dirBaseName+['W'+self.nomega]
        #if self.nbandsgw:
        #    self.parameter=self.parameter+[' -B '+self.nbandsgw]
        #else:
        #    self.parameter=self.parameter+[' ']
        #self.dirBaseName=self.dirBaseName+['G'+self.nbandsgw]
        if self.encutgw:
            self.parameter=self.parameter+[' -E '+self.encutgw]
        else:
            self.parameter=self.parameter+[' ']
        self.dirBaseName=self.dirBaseName+['E'+self.encutgw]   
        if self.encut:
            self.parameter=self.parameter+[' -e '+self.encut]
        else:
            self.parameter=self.parameter+[' ']
        self.dirBaseName=self.dirBaseName+['e'+self.encut]  
        if self.prec:
            self.parameter=self.parameter+[' -p '+self.prec]
            self.dirBaseName=self.dirBaseName+['p'+self.prec[0]]  
        else:
            self.parameter=self.parameter+[' ']
            self.dirBaseName=self.dirBaseName+[''] 
        
        suffix='_'+self.functional+'GW.VASP'
        if self.steps=='0':
            steps=' -j PBE;'+self.functional+';G0W0 -i 0'
            suffix='_'+self.functional+'GW.VASP'
        if self.steps=='1':
            steps=' -j '+self.functional+';G0W0'
            suffix='_'+self.functional+'GW.VASP'
        if self.steps=='2':
            steps=' -j PBE;G0W0 -i 0'
            suffix='_PBEGW.VASP'
        if self.steps=='3':
            steps=' -j PBE;'+self.functional+' -i 0'
            suffix='_'+self.functional+'.VASP'
        if self.steps=='5':
            steps=' -j '+self.functional
            suffix='_'+self.functional+'.VASP'
        if self.steps=='4':
            steps=' -j PBE -i 0'
            suffix='_PBE.VASP'
        if self.steps=='6':
            steps=' -j G0W0'
            suffix='_'+self.functional+'GW.VASP'
        #print self.parameter
        #print self.dirBaseName 
                            
        if self.kplist:
            k1f=int(args.kpoints.split(';')[0].split(',')[0]) 
            k1l=int(args.kpoints.split(';')[1].split(',')[0])
            try:
                k1s=int(args.kpoints.split(';')[2].split(',')[0])
            except IndexError:
                k1s=2
            if len(args.kpoints.split(';')[0].split(','))>1:
                k2f=int(args.kpoints.split(';')[0].split(',')[1])
                k2l=int(args.kpoints.split(';')[1].split(',')[1])
                try:k2s=int(args.kpoints.split(';')[2].split(',')[1])
                except IndexError:k2s=2
                if len(args.kpoints.split(';')[0].split(','))==3:
                    k3f=int(args.kpoints.split(';')[0].split(',')[2])
                    k3l=int(args.kpoints.split(';')[1].split(',')[2])
                    try:k3s=int(args.kpoints.split(';')[2].split(',')[2])
                    except IndexError:k3s=2
                else:
                    k3f,k3l,k3s=1,1,1
            else:
                k2f,k2l,k2s=1,1,1
                k3f,k3l,k3s=1,1,1
            k1=range(k1f,k1l+k1s,k1s)
            k2=range(k2f,k2l+k2s,k2s)
            k3=range(k3f,k3l+k3s,k3s)
            if len(k2)==1:
                k2=k2*len(k1)
            elif len(k2) != len(k1):
                print "number of k2 "+str(len(k2))+" do not equal to number of k1 "+str(len(k1))
                sys.exit()
            if len(k3)==1:
                k3=k3*len(k1)
            elif len(k3) != len(k1):
                print "number of k2 "+len(3)+" do not equal to number of k1 "+len(k1)
                sys.exit()
            k1w=len(str(k1l))
            k2w=len(str(k2l))
            k3w=len(str(k3l))
            klist=[','.join([str(k1[i]).zfill(k1w),str(k2[i]).zfill(k2w),str(k3[i]).zfill(k3w)]) for i in range(len(k1))]
            dirSuffix=''.join(self.dirBaseName[1:])+suffix
            paramSuffix=''.join(self.parameter[1:])+' -n '+self.nproc+' -j PBE;'+self.functional+';G0W0 -i 0'
            if self.steps and self.steps in '012345':
                paramSuffix=''.join(self.parameter[1:])+' -n '+self.nproc+steps

                
            #print param
            #print dirSuffix
            for kp in klist:
                dirName='K'+''.join(kp.split(','))+'_'+dirSuffix
                parameter='inpgen_vasp.py -k '+kp+' '+paramSuffix
                print parameter.split()
                try:
                    shutil.copytree(self.baseDIR,dirName)
                except OSError:
                    pass
                os.chdir(dirName)
                subprocess.Popen(parameter.split())
                os.chdir(self.currentDIR)
                
        if self.nblist:
            try:
                nbf,nbl,nbs=map(int,self.nblist.split(';'))
            except ValueError:
                print "Error, must input 3 number to set the range and seperate by semicolon"
                sys.exit()
            nbw=len(str(nbl))
            nblist=range(nbf,nbl+nbs,nbs)
            dirSuffix=''.join([self.dirBaseName[0]]+self.dirBaseName[2:])+suffix
            paramSuffix=''.join([self.parameter[0]]+self.parameter[2:])+' -n '+self.nproc+' -j '+self.functional+';G0W0 '
            if self.steps and  self.steps in '012345':
                paramSuffix=''.join([self.parameter[0]]+self.parameter[2:])+' -n '+self.nproc+steps
            for nb in nblist:
                dirName='B'+str(nb).zfill(nbw)+'_'+dirSuffix
                parameter='inpgen_vasp.py -b '+str(nb)+' '+paramSuffix
                print parameter.split()
                try:
                    shutil.copytree(self.baseDIR,dirName)
                except OSError:
                    pass
                os.chdir(dirName)
                subprocess.Popen(parameter.split())
                os.chdir(self.currentDIR)
 
                        
        if self.nwlist:
            try:
                nwf,nwl,nws=map(int,self.nwlist.split(';'))
            except ValueError:
                print "Error, must input 3 number to set the range and seperate by semicolon"
                sys.exit() 
            nww=len(str(nwl))
            nwlist=range(nwf,nwl+nws,nws)
            dirSuffix=''.join(self.dirBaseName[:2]+self.dirBaseName[3:])+suffix
            paramSuffix=''.join(self.parameter[:2]+self.parameter[3:])+' -n '+self.nproc+' -j G0W0 '
            if self.steps and self.steps in '012':
                paramSuffix=''.join(self.parameter[:2]+self.parameter[3:])+' -n '+self.nproc+steps
            for nw in nwlist:
                dirName='W'+str(nw).zfill(nww)+'_'+dirSuffix
                parameter='inpgen_vasp.py -w '+str(nw)+' '+paramSuffix
                print parameter.split()
                try:
                    shutil.copytree(self.baseDIR,dirName)
                except OSError:
                    pass
                os.chdir(dirName)
                subprocess.Popen(parameter.split())
                os.chdir(self.currentDIR)
        
        if self.eglist:
            try:
                egf,egl,egs=map(int,self.eglist.split(';'))
            except ValueError:
                print "Error, must input 3 number to set the range and seperate by semicolon"
                sys.exit() 
            egw=len(str(egl))
            eglist=range(egf,egl+egs,egs)
            dirSuffix=''.join(self.dirBaseName[:3]+self.dirBaseName[4:])+suffix
            paramSuffix=''.join(self.parameter[:3]+self.parameter[4:])+' -n '+self.nproc+' -j G0W0 '
            if self.steps and self.steps in '012':
                paramSuffix=''.join(self.parameter[:3]+self.parameter[4:])+' -n '+self.nproc+steps
            for eg in eglist:
                dirName='E'+str(eg).zfill(egw)+'_'+dirSuffix
                parameter='inpgen_vasp.py -E '+str(eg)+' '+paramSuffix
                print parameter.split()
                try:
                    shutil.copytree(self.baseDIR,dirName)
                except OSError:
                    pass
                os.chdir(dirName)
                subprocess.Popen(parameter.split())
                os.chdir(self.currentDIR)
                
        #if self.nglist:
        #    try:
        #        ngf,ngl,ngs=map(int,self.nglist.split(';'))
        #    except ValueError:
        #        print "Error, must input 3 number to set the range and seperate by semicolon"
        #        sys.exit() 
        #    ngw=len(str(ngl))
        #    nglist=range(ngf,ngl,ngs)
        #    dirSuffix=''.join(self.dirBaseName[:3]+self.dirBaseName[4:])+'.VASP'
        #    paramSuffix=''.join(self.parameter[:3]+self.parameter[4:])+' -n '+self.nproc+' -j G0W0 '
        #    for ng in nglist:
        #        dirName='G'+str(ng).zfill(ngw)+'_'+dirSuffix
        #        parameter='inpgen_vasp.py -B '+str(ng)+' '+paramSuffix
        #        print parameter.split()
        #        shutil.copytree(self.baseDIR,dirName)
        #        os.chdir(dirName)
        #        subprocess.Popen(parameter.split())
        #        os.chdir(self.currentDIR)
                    
        if self.eclist:
            try:
                ecf,ecl,ecs=map(int,self.eclist.split(';'))
            except ValueError:
                print "Error, must input 3 number to set the range and seperate by semicolon"
                sys.exit()
            ecw=len(str(ecl))
            eclist=range(ecf,ecl+ecs,ecs)
            dirSuffix=''.join(self.dirBaseName[:4]+self.dirBaseName[5:])+suffix
            paramSuffix=''.join(self.parameter[:4]+self.parameter[5:])+' -n '+self.nproc+' -j PBE;'+self.functional+';G0W0 '
            if self.steps and self.steps in '0235':
                paramSuffix=''.join(self.parameter[:4]+self.parameter[5:])+' -n '+self.nproc+steps
            
            for ec in eclist:
                dirName='e'+str(ec).zfill(ecw)+'_'+dirSuffix
                parameter='inpgen_vasp.py -e '+str(ec)+' '+paramSuffix
                print parameter.split()
                try:
                    shutil.copytree(self.baseDIR,dirName)
                except OSError:
                    pass
                os.chdir(dirName)
                subprocess.Popen(parameter.split())
                os.chdir(self.currentDIR)       
            
        
            
        
    
parser=argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter,description="A script used with Gaussian09 to auto optimize omega and alpha for tuning DFT.")
parser.add_argument("-k","--kpoints",help="Set range to test kpoints. Input 3 group for start,end and step, e.g. 2,1,1;16,1,1;2,0,0",  type=str,metavar='KPOINTS range')
parser.add_argument("-K","--KPOINTS",help="Set a single K point.  Default: none",  type=str,default='')  
#parser.add_argument("-g","--nbandsgw",help="Set list of NBANDSGW. Input 3 number for start end and step, e.g. -n 100;200;20", type=str,metavar='NBANDSGW range')
#parser.add_argument("-G","--NBANDSGW",help="Set the nbandsgw", type=str)
parser.add_argument("-b","--nbands",help="Set list of nbands. Input 3 number for start end and step, e.g. -b 100;200;20", type=str,metavar='NBANDS range') 
parser.add_argument("-B","--NBANDS",help="Set the NBANDS",  type=str)  
parser.add_argument("-e","--encutgw",help="Set list of ENCUTGW. Input 3 number for start end and step, e.g. -e 100;200;20", type=str,metavar='ENCUTGW range')
parser.add_argument("-E","--ENCUTGW",help="Set the ENCUTGW", type=str)
parser.add_argument("-c","--encut",help="Set list of ENCUT. Input 3 number for start end and step, e.g. -e 200;400;20", type=str,metavar='ENCUT range')
parser.add_argument("-C","--ENCUT",help="Set the ENCUT", type=str)
#parser.add_argument("-t","--tuning",help="Set +the tuning scheme, available options are: w(omega),a(alpha),aw(alpha and omega),s(scan alpha or omege, must be used with -a or -w option)\nPut m in front of above options (mw,ma,maw,ms) means tuning with multiple structure,\n\
#e.g. \"-t mw\" means optimize the omega to minimize sum of the criterion of all the input structure.\nIf different criterion is used for different structure, use -c option to designate criterion explicitly for each structure, \ne.g. \"-t ma -c Jh,Jl\" means optimize alpha to minimize sum of Jh of structure 1 and Jl of structure 2.  Default: w", type=str,default='w') 
parser.add_argument("-w","--nomega",help="Set list of NOMEGA. Input 3 number for start end and step, e.g. -w 20;50;5", type=str,metavar='NOMEGA range')
parser.add_argument("-W","--NOMEGA",help="Set the value of NOMEGA. Default: None", type=str) 
parser.add_argument("-f","--functional",help="Set the base functional. Default: HSE03", type=str,default='HSE03') 
parser.add_argument("-s","--steps",help="calculation steps,0=pbe+hybrid+gw,1=hybrid+gw,2=pbe+gw,3=pbe+hybrid,4=pbe,5=hybrid,6=gw default:gw + smallst steps needed",type=str,default='')
parser.add_argument("-n","--nproc",help="Set the number of processors. Default: 16", type=str,default='16') 
parser.add_argument("baseD",nargs=1,help="must input one base directory contain POSCAR and POTCAR ")
args=parser.parse_args()


baseD=args.baseD[0]
if not os.path.exists(baseD):
    print "Error!!! Target Dierctory "+baseD+" do not exist!"
    sys.exit()
conver=converge()
conver.baseDIR=os.path.abspath(baseD)
conver.readparam()
conver.nproc=args.nproc
conver.functional=args.functional

if args.kpoints:
    conver.kplist=args.kpoints
if args.nbands:
    conver.nblist=args.nbands
if args.encutgw:
    conver.eglist=args.encutgw
if args.encut:
    conver.eclist=args.encut
#if args.nbandsgw:
#    conver.nglist=args.nbandsgw
if args.nomega:
    conver.nwlist=args.nomega

if args.NOMEGA:
    conver.nomega=args.NOMEGA
if args.NBANDS:
    conver.nbands=args.NBANDS
    conver.steps='1'
if args.KPOINTS:
   conver.kpoint=args.KPOINTS
   conver.steps='0'
#if args.NBANDSGW:
#    conver.nbandsgw=args.NBANDSGW
if args.ENCUTGW:
    conver.encutgw=args.ENCUTGW
if args.ENCUT:
    conver.encut=args.ENCUT
    conver.steps='0'
if args.steps:
    conver.steps=args.steps
conver.setformation()

    #for kp in klist:
    #    targetDIR='K'+''.join(kp.split(','))+'_base'
    #    shutil.copytree(baseD,targetDIR)
        
        
    